import torch
from torch import nn


class GanNetwork(nn.Module):

    def __init__(self, optimizer=torch.optim.SGD, learningSpeed=1e-2):
        super(GanNetwork, self).__init__()
        self.optimizer = optimizer
        self.learningSpeed = learningSpeed
        self.discriminator = nn.Sequential(
            nn.Linear(28 * 28, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 1),
            nn.Sigmoid()
        )

        self.generator = nn.Sequential(
            nn.Linear(3, 128),
            nn.ReLU(),
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, 28 * 28)
        )

        self.discriminator_optimizer = self.optimizer(self.discriminator.parameters(), lr=self.learningSpeed)
        self.generator_optimizer = self.optimizer(self.generator.parameters(), lr=self.learningSpeed)

    def forward(self):
        raise NotImplemented("此方法在这个网络中定义不明，所以不实现")

    def train(self, seed, x):
        generated_data = self.generator.forward(seed)

        generated_data_output = self.discriminator(generated_data);
        real_data_output = self.discriminator(x)

        generated_loss = torch.mean(torch.log(1 - generated_data_output))
        real_loss = - torch.mean(torch.log(real_data_output) + torch.log(1 - generated_data_output))

        self.discriminator_optimizer.zero_grad()
        real_loss.backward(retain_graph=True)
        self.discriminator_optimizer.step()
        self.generator_optimizer.zero_grad()
        generated_loss.backward()
        self.generator_optimizer.step()
        return generated_data
